# encoding=utf-8
"""
需要这两个环境变量, 在pycharm命令行启动前设置
source /usr/local/Ascend/ascend-toolkit/set_env.sh
export PYTHONPATH=/usr/local/Ascend/thirdpart/aarch64/acllite:$PYTHONPATH
"""
import time
import cv2
import numpy
from AclTools import AclLiteResource
from acllite_model import AclLiteModel


class OCTDeepLabInfer:
    def __init__(self, pth_path: str):
        self.acl_resource = AclLiteResource()
        self.acl_resource.init()
        self.model = AclLiteModel(pth_path)

        self.colors = {
            1: (0, 0, 255),  # Red
            2: (0, 255, 0),  # Green
            3: (255, 0, 0),  # Blue
            4: (0, 255, 255),  # Yellow
            5: (255, 0, 255)  # Magenta
        }

    def infer(self, image_path: str):
        trans_dim, ori_size = self._preprocess(image_path)
        result = self.model.execute([trans_dim, ])
        return self._postprocess(result[0], ori_size)

    def _postprocess(self, pred_arr: numpy.ndarray, ori_size: tuple):
        decode_array = numpy.argmax(pred_arr, axis=1)[0]
        H, W = ori_size
        img_arr = cv2.resize(decode_array.astype(numpy.uint8), (W, H), interpolation=cv2.INTER_NEAREST)

        masks = numpy.zeros((H, W, 3), dtype=numpy.uint8)
        for index, color in self.colors.items():
            masks[img_arr == index] = color

        return img_arr, masks

    def _preprocess(self, data_in: str):
        # load
        ori_img = cv2.imread(data_in)
        ori_size = ori_img.shape[:2]  # h, w

        # transform
        trans_image = cv2.cvtColor(ori_img, cv2.COLOR_BGR2RGB)
        trans_image = cv2.resize(trans_image, (512, 512), interpolation=cv2.INTER_CUBIC)
        trans_image = trans_image / 255.0

        # return
        x = numpy.expand_dims(trans_image.astype(numpy.float32).transpose((2, 0, 1)), axis=0)
        return x, ori_size

if __name__ == '__main__':
    # init
    t1 = time.perf_counter()
    seg = OCTDeepLabInfer("oct_om_seg.om")
    t2 = time.perf_counter()
    print("init cost", t2 - t1)

    # time cost
    t1 = time.perf_counter()
    ori_array, color_array = seg.infer("sample/ori.jpg")
    t2 = time.perf_counter()
    print("infer cost", t2 - t1)

    # vis
    cv2.imwrite("sample/raw.png", ori_array)
    cv2.imwrite("sample/color.png", color_array)

    # merge
    data1 = cv2.imread("sample/ori.jpg")
    data2 = cv2.imread("sample/color.png")
    data3 = cv2.addWeighted(data1, 0.7, data2, 0.3, 3)
    cv2.imwrite("sample/merge.png", data3)